import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
from flappybird.game.wrapped_flappy_bird import GameState
from util import convert_image
# import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('D:/PythonProject/tensorboarddir')

# Hyperparameters
learning_rate = 0.005
gamma = 0.98
lmbda = 0.95
eps_clip = 0.1
K_epoch = 5  # 一笔data练几次
T_horizon = 32 # 一次取多少笔data

device = 'cuda:0'

class PPO(nn.Module):
    def __init__(self):
        super(PPO, self).__init__()
        self.data = []
        self.gstep = 0

        self.cnn = nn.Sequential(
            # conv1层，输入的灰度图，所以 in_channels=1, out_channels=8 说明使用了8个滤波器/卷积核，
            # kernel_size=3卷积核大小3x3
            # input 維度 [1, 80, 80]
            nn.Conv2d(1, 8, 3, 1, 1),  # [8, 80, 80] 这边的注释的是输出的维度
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),  # [8, 40, 40]

            nn.Conv2d(8, 16, 3, 1, 1),  # [16, 40, 40]
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),  # [16, 20, 20]

            nn.Conv2d(16, 16, 3, 1, 1),  # [16, 20, 20]
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2, 2, 0),  # [16, 10, 10]
        )

        self.fc = nn.Sequential(  # 定义全连接层，使用Linear激活函数
            nn.Linear(1600, 800),
            nn.ReLU(),
            nn.Linear(800, 256),
        )

        self.fc_pi = nn.Linear(256, 2)
        self.fc_v = nn.Linear(256, 1)
        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)

    def pi(self, x, softmax_dim=0):
        x = x.to(device)
        x = self.cnn(x)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc(x))
        x = self.fc_pi(x)
        prob = F.softmax(x, dim=softmax_dim)
        return prob

    def v(self, x):
        x = x.to(device)
        x = self.cnn(x)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc(x))
        v = self.fc_v(x)
        return v

    def put_data(self, transition):
        self.data.append(transition)

    def make_batch(self):
        s_lst, a_lst, r_lst, s_prime_lst, prob_a_lst, done_lst = [], [], [], [], [], []
        for transition in self.data:
            s, a, r, s_prime, prob_a, done = transition

            s_lst.append(s)
            a_lst.append([a])
            r_lst.append([r])
            s_prime_lst.append(s_prime)
            prob_a_lst.append([prob_a])
            done_mask = 0 if done else 1
            done_lst.append([done_mask])

        s, a, r, s_prime, done_mask, prob_a = torch.stack(s_lst, dim=0), torch.tensor(a_lst), \
                                              torch.tensor(r_lst), torch.stack(s_prime_lst, dim=0), \
                                              torch.tensor(done_lst, dtype=torch.float), torch.tensor(prob_a_lst)
        self.data = []
        return s, a, r, s_prime, done_mask, prob_a

    def train_net(self):
        s, a, r, s_prime, done_mask, prob_a = self.make_batch()

        # s = s.to(device)
        # s_prime = s_prime.to(device)

        for i in range(K_epoch):
            td_target = r.to(device) + gamma * self.v(s_prime) * done_mask.to(device)
            delta = td_target - self.v(s)
            delta = delta.detach().cpu().numpy()

            advantage_lst = []
            advantage = 0.0
            for delta_t in delta[::-1]:
                advantage = gamma * lmbda * advantage + delta_t[0]
                advantage_lst.append([advantage])
            advantage_lst.reverse()
            advantage = torch.tensor(advantage_lst, dtype=torch.float).to(device)

            pi = self.pi(s, softmax_dim=1)
            pi_a = pi.gather(1, a.to(device))
            ratio = torch.exp(torch.log(pi_a) - torch.log(prob_a.to(device)))  # a/b == exp(log(a)-log(b))

            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - eps_clip, 1 + eps_clip) * advantage
            loss = -torch.min(surr1, surr2) + F.smooth_l1_loss(self.v(s), td_target.detach())

            self.optimizer.zero_grad()
            loss.mean().backward()
            self.optimizer.step()
            writer.add_scalar('avg_loss-ppo', torch.sum(loss) / T_horizon, self.gstep)
            self.gstep+=1

def main():
    # env = gym.make('CartPole-v1')
    env = GameState()
    # torch.cuda.set_device(0)
    model = PPO().cuda()
    score = 0.0
    print_interval = 50

    for n_epi in range(10000):
        s = env.reset()
        s = convert_image(s)
        done = False
        while not done:
            for t in range(T_horizon):
                prob = model.pi(s.unsqueeze(0))
                m = Categorical(prob)
                a = m.sample().item()
                s_prime, r, done, info = env.step(a)
                s_prime = convert_image(s_prime)

                model.put_data((s, a, r / 2.0, s_prime, prob[0][a].item(), done))
                s = s_prime

                score += r
                if done:
                    break

            model.train_net()

        if n_epi % print_interval == 0 and n_epi != 0:
            print("episode :{}, avg score : {:.1f}".format(n_epi, score / print_interval))
            score = 0.0
            torch.save(model.state_dict(), 'saved_networks/ppo-flappybird.pkl')

    # env.close()


if __name__ == '__main__':
    main()